package drds.plus.datanode.select.equity_manager_select;

import drds.plus.datanode.configuration.DataNodeConfigurationManager;
import drds.plus.datanode.configuration.DataNodeExtraConfiguration;
import drds.plus.datanode.configuration.DataSourceWrapper;
import drds.plus.datanode.executor.Executor;
import drds.plus.datanode.select.AbstractSelector;
import drds.plus.datanode.select.DataSourceHolder;
import drds.plus.datasource.api.DataSource;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import java.sql.SQLException;
import java.util.*;

/**
 * 对等数据库管理器 可以是读对等：如多个读库，每个库的数据完全相同。对等读取 可以是写对等：如日志库，每个库数据不同，一条数据写入哪个库都可以。对等写入
 * 支持动态推送权重，动态加减库
 */

// 因为当配置信息变动时每次都会重新生成一个新的EquityDbManager实例，
// 所以原有的与"动态改变"相关的代码在新的EquityDbManager实现中已全部删除
@Slf4j
public class EquityManager extends AbstractSelector {
    private final Random random = new Random();
    @Setter
    @Getter
    private Map<String, DataSourceHolder> dataSourceIdToDataSourceHolderMap;
    @Setter
    @Getter
    private RandomWeightSelect randomWeightSelect;

    public EquityManager(Map<String, DataSourceWrapper> dataSourceWrapperMap, Map<String, Integer> weightMap) {
        this.dataSourceIdToDataSourceHolderMap = new HashMap<String, DataSourceHolder>(dataSourceWrapperMap.size());
        for (Map.Entry<String, DataSourceWrapper> e : dataSourceWrapperMap.entrySet()) {
            this.dataSourceIdToDataSourceHolderMap.put(e.getKey(), new DataSourceHolder(e.getValue()));
        }
        this.randomWeightSelect = new RandomWeightSelect(weightMap);
    }

    public EquityManager(Map<String, DataSourceWrapper> dataSourceWrapperMap, Map<String, Integer> weightMap, DataNodeExtraConfiguration dataNodeExtraConfiguration) {
        super.dataNodeExtraConfiguration = dataNodeExtraConfiguration;
        this.dataSourceIdToDataSourceHolderMap = new HashMap<String, DataSourceHolder>(dataSourceWrapperMap.size());
        for (Map.Entry<String, DataSourceWrapper> e : dataSourceWrapperMap.entrySet()) {
            this.dataSourceIdToDataSourceHolderMap.put(e.getKey(), new DataSourceHolder(e.getValue()));
        }
        this.randomWeightSelect = new RandomWeightSelect(weightMap);
    }

    private static String selectAliveKey(RandomWeightSelect randomWeightSelect, List<String> excludeKeys) {
        if (null == excludeKeys) {
            excludeKeys = new ArrayList<String>();
        }
        return randomWeightSelect.select(excludeKeys);
    }

    /**
     * @return 根据权重，随机返回一个DataSource
     */
    public DataSourceWrapper select() {
        String aliveKey = selectAliveKey(randomWeightSelect, null);
        if (null != aliveKey) {
            return this.get(aliveKey);
        } else {
            return null;
        }
    }

    /**
     * 返回指定dsKey对应的数据源。若对应数据源的当前权重为0，则返回null
     *
     * @param dataSourceId 内部和每一个物理DataSource对应的key, 在初始化dbSelector时指定
     * @return 返回dsKey对应的数据源
     */
    public DataSourceWrapper get(String dataSourceId) {
        DataSourceHolder dataSourceHolder = dataSourceIdToDataSourceHolderMap.get(dataSourceId);
        Integer weigthValue = this.randomWeightSelect.getWeightConfigMap().get(dataSourceId);
        if (weigthValue == null || weigthValue.equals(0)) {
            return null;
        }
        return dataSourceHolder == null ? null : dataSourceHolder.dataSourceWrapper;
    }

    // TODO 考虑接口是否缩小为只返回DataSource[]
    public Map<String, DataSource> getDataSourceIdToDataSourceHolderMap() {
        Map<String, DataSource> dataSourceIdToDataSourceMap = new HashMap<String, DataSource>(this.dataSourceIdToDataSourceHolderMap.size());
        for (Map.Entry<String, DataSourceHolder> e : this.dataSourceIdToDataSourceHolderMap.entrySet()) {
            dataSourceIdToDataSourceMap.put(e.getKey(), e.getValue().dataSourceWrapper.getDataSource());
        }
        return dataSourceIdToDataSourceMap;
    }


    /**
     * 在所管理的数据库上重试执行一个回调操作。失败了根据权重选下一个库重试 以根据权重选择到的DataSource，和用户传入的自用参数args，
     * 重试调用DataSourceTryer的tryOnDataSource方法
     *
     * @param failedDataSourceToSQLExceptionMap 已知的失败DS及其异常
     * @param args                              透传到DataSourceTryer的tryOnDataSource方法中
     * @return null表示执行成功。否则表示重试次内执行失败，返回SQLException列表
     */
    public <T> T selectDataSourceHolderAndGetConnectionWrapperAndExecuteAndRetryIfFailed(Map<javax.sql.DataSource, SQLException> failedDataSourceToSQLExceptionMap, Executor<T> executor, int times, Object... args) throws SQLException {
        // 如果不支持重试，把times设为1就可以了
        if (!this.isSupportRetry()) {
            times = 1;
        }
        //
        RandomWeightSelect randomWeightSelect = this.randomWeightSelect;
        List<SQLException> sqlExceptionList = new ArrayList<SQLException>(0);
        List<String> excludeKeys = new ArrayList<String>(0);
        if (failedDataSourceToSQLExceptionMap != null) {
            sqlExceptionList.addAll(failedDataSourceToSQLExceptionMap.values());
            times = times - failedDataSourceToSQLExceptionMap.size(); // 扣除已经失败掉的重试次数
            //
            for (SQLException e : failedDataSourceToSQLExceptionMap.values()) {
                if (!exceptionSorter.isFatalException(e)) {
                    // 有一个异常（其实是最后加入的异常，因map无法知道顺序，只能遍历）不是数据库不可用异常，则抛出 是不是应该在发现非数据库fatal之后就立刻抛出，而不是放到failedDataSources这个map里?(guangxia)
                    return executor.onSqlException(sqlExceptionList, exceptionSorter, args);
                }
            }
        }
        for (int i = 0; i < times; i++) {
            String aliveKey = selectAliveKey(randomWeightSelect, excludeKeys);
            if (aliveKey == null) {
                sqlExceptionList.add(new GroupNotAvaliableException("tryTime:" + i + ", excludeKeys:" + excludeKeys + ", weightConfig:" + randomWeightSelect.getWeightConfigMap()));
                break;//异常了 直接结束
            }
            DataSourceHolder dataSourceHolder = dataSourceIdToDataSourceHolderMap.get(aliveKey);
            if (dataSourceHolder == null) {
                // 不应该出现的。初始化逻辑应该保证空的数据源(null)不会被加入dataSourceMap
                throw new IllegalStateException("Can't find DataSource for id:" + aliveKey);
            }
            if (failedDataSourceToSQLExceptionMap != null && failedDataSourceToSQLExceptionMap.containsKey(dataSourceHolder.dataSourceWrapper)) {
                excludeKeys.add(aliveKey);
                i--; // 这次不算重试次数
                continue;
            }
            // TODO 有必要每次都检查DataSource的状态吗 检查一下数据源，如果是NA或往一个只读的库中写记录都要重试下一个数据源
            if (!DataNodeConfigurationManager.isDataSourceAvailable(dataSourceHolder.dataSourceWrapper, this.readable)) {
                excludeKeys.add(aliveKey);
                i--; // 这次不算重试次数
                continue;
            }
            try {
                if (dataSourceHolder.isNotAvailable) {
                    boolean needRetry = System.currentTimeMillis() - dataSourceHolder.lastRetryTime > retryBadDbInterval;
                    if (needRetry && dataSourceHolder.lock.tryLock()) {
                        try {
                            T t = executor.createConnectionWrapperAndExecute(dataSourceHolder.dataSourceWrapper, args); // 同一个时间只会有一个线程继续使用这个数据源。
                            dataSourceHolder.isNotAvailable = false; // 用一个线程重试，执行成功则标记为可用，自动恢复
                            return t;
                        } finally {
                            dataSourceHolder.lastRetryTime = System.currentTimeMillis();
                            dataSourceHolder.lock.unlock();
                        }
                    } else {
                        excludeKeys.add(aliveKey); // 其他线程跳过已经标记为notAvailable的数据源
                        i--; // 这次不算重试次数
                        continue;
                    }
                } else {
                    return executor.createConnectionWrapperAndExecute(dataSourceHolder.dataSourceWrapper, args); // 有一次成功直接返回
                }
            } catch (SQLException e) {
                sqlExceptionList.add(e);
                boolean fatalException = exceptionSorter.isFatalException(e);
                if (fatalException) {
                    dataSourceHolder.isNotAvailable = true;
                }
                if (!fatalException || failedDataSourceToSQLExceptionMap == null) {
                    // throw e; //如果不是数据库不可用异常，或者不要求重试，直接抛出
                    break;
                }
                log.warn(new StringBuilder().append(i + 1).append("th try locate on [").append(aliveKey).append("] failed:").append(e.getMessage()).toString()); // 这里不打异常栈了,全部重试失败才由调用者打
                excludeKeys.add(aliveKey);
            }
        }
        return executor.onSqlException(sqlExceptionList, exceptionSorter, args);
    }

    /**
     * <pre>
     * 分流：随机返回权重串里包含值为dataSourceIndex的i的数据源
     * 如果权重串没有定义i/I，则dataSourceIndex等于几，就路由到group中的第几个的数据源
     * 一个db可以同时配置多个i；不同的db可以配置相同的i，
     * 例如权重串= db0:rwi0i2, db1:ri1, db2:ri1, db3:ri2
     * a. 用户指定dataSourceIndex=0，路由到db0；（只有db0有i0）
     * b. 用户指定dataSourceIndex=1，随机路由到db1和db2；（db1和db2都有i1）
     * c. 用户指定dataSourceIndex=2，随机路由到db0和db3；（db0和db3都有i2）
     * d. 如果没有配置i，例如db0:rw, db1:readWeight;指定dataSourceIndex=1则路由到db1
     * </pre>
     */
    public DataSourceHolder findDataSourceHolderByDataSourceIndex(String dataSourceIndex) {
        List<DataSourceHolder> dataSourceHolderList = new ArrayList<DataSourceHolder>();
        for (DataSourceHolder dataSourceHolder : dataSourceIdToDataSourceHolderMap.values()) {
            if (dataSourceHolder.dataSourceWrapper.isMatchDataSourceIndex(dataSourceIndex))
                dataSourceHolderList.add(dataSourceHolder);
        }
        if (!dataSourceHolderList.isEmpty()) {
            return dataSourceHolderList.get(random.nextInt(dataSourceHolderList.size()));
        } else {
            return null;
        }
    }
}
